#ifndef _UTILS_ALGORITHM_H_
#define _UTILS_ALGORITHM_H_

#include <iostream>
#include <vector>
#include <queue>
#include <limits>
#include <map>
#include <stack>
#include <cmath>
#include <algorithm>

using namespace std;

class TopFinder
{
public:
  struct ValIndex
  {
    double value;
    int index;
  };

  void findTopNWithMinIndexGap(const std::vector<ValIndex> &sortedArr,
                               int n, int a,
                               std::vector<ValIndex> &result, int start)
  {
    if (n == 0) return;

    for (int i = start; i < sortedArr.size(); ++i) {
      bool isFarEnough = true;
      for (const auto &res : result) {
        if (abs(sortedArr[i].index - res.index) < a) {
          isFarEnough = false;
          break;
        }
      }
      if (isFarEnough) {
        result.push_back(sortedArr[i]);
        if (result.size() == n) return;
        findTopNWithMinIndexGap(sortedArr, n, a, result, i + 1);
        if (result.size() == n) return;
      }
    }
  }

  /**
   * @brief find top n elements with minimum index gap
   *
   * @param arr input array
   * @param n top n
   * @param a minimum index gap
   * @return std::vector<int> top n indices
   */
  std::vector<int> findTopN(const std::vector<double> &arr,
                            int n, int a)
  {
    std::vector<ValIndex> valIndexArr;
    for (int i = 0; i < arr.size(); ++i) {
      valIndexArr.push_back({arr[i], i});
    }

    // Sort in descending order based on value
    std::sort(valIndexArr.begin(), valIndexArr.end(),
              [](const ValIndex &a, const ValIndex &b) {
                return a.value > b.value;
              });

    std::vector<ValIndex> result;
    findTopNWithMinIndexGap(valIndexArr, n, a, result, 0);

    // Extract values from result
    std::vector<int> finalResult;
    for (const auto &res : result) {
      finalResult.push_back(res.index);
    }

    return finalResult;
  }
};

class GimbalOptimizer
{
private:
  struct AngleGain
  {
    double angle;
    double gain;
  };

  map<double, vector<AngleGain>> angleOptions;
  vector<double> times; // Store the times in sequence
  double maxAngleSpeed{1.57};
  double gazeTime{0.0};

  // To store the optimal time-angle sequence
  vector<pair<double, double>> optimalSequence;

  struct Node
  {
    double time;
    double angle;
    double cost;
    Node(double t, double a, double c) : time(t), angle(a), cost(c) {}
    bool operator>(const Node &other) const
    {
      return cost > other.cost;
    }
  };

public:
  GimbalOptimizer() {}
  GimbalOptimizer(double maxSpeed) : maxAngleSpeed(maxSpeed) {}

  void setLimits(double _maxSpeed, double _gazeTime)
  {
    maxAngleSpeed = _maxSpeed;
    gazeTime = _gazeTime;
  }

  bool setAngleOptions(const vector<double> &_times,
                       const vector<vector<double>> &_angles,
                       const vector<vector<double>> &_gains)
  {
    // check if the input is valid
    size_t n = _times.size();
    if (n != _angles.size() || n != _gains.size())
    {
      std::cerr << "Invalid input: the number of times,"
                << " angles, and gains should be the same." << endl;
      return false;
    }
    for (size_t i = 0; i < n; ++i)
    {
      if (_angles[i].size() != _gains[i].size())
      {
        std::cerr << "Invalid input: the number of angles"
                  << " and gains should be the same." << endl;
        return false;
      }
    }

    times = _times;
    angleOptions.clear();
    for (size_t i = 0; i < n; ++i) {
      vector<AngleGain> angleGainPairs(_angles[i].size());
      for (size_t j = 0; j < _angles[i].size(); ++j) {
        angleGainPairs[j] = {_angles[i][j], _gains[i][j]};
      }
      angleOptions[_times[i]] = angleGainPairs;
    }
    return true;
  }

  bool canTransition(double deltaAngle, double timeInterval)
  {
    return abs(deltaAngle) <= timeInterval * maxAngleSpeed + 0.001;
  }

  bool optimize(double initialTime, double initialAngle)
  {
    if (initialTime > times.front()) {
      cerr << "Initial time is out of range." << endl;
      return false;
    }

    times.insert(times.begin(), initialTime);

    priority_queue<Node, vector<Node>, greater<Node>> pq;
    map<pair<double, double>, double> dist;
    map<pair<double, double>, pair<double, double>> prev;

    pq.push(Node(initialTime, initialAngle, 0.0));
    dist[{initialTime, initialAngle}] = 0.0;

    while (!pq.empty()) {
      Node current = pq.top();
      pq.pop();

      std::vector<double>::iterator currentTimeIt =
          find(times.begin(), times.end(), current.time);
      if (currentTimeIt == times.end()) continue;
      if (next(currentTimeIt) == times.end()) {
        stack<pair<double, double>> path;
        pair<double, double> node = {current.time, current.angle};

        while (prev.find(node) != prev.end()) {
          path.push(node);
          node = prev[node];
        }
        path.push({initialTime, initialAngle});

        optimalSequence.clear();
        while (!path.empty()) {
          optimalSequence.push_back(path.top());
          path.pop();
        }

        return true;
      }

      double nextTime = *next(currentTimeIt);
      double timeInterval = nextTime - current.time;
      if (current.time != initialTime) timeInterval -= gazeTime;

      for (const AngleGain &ag : angleOptions[nextTime]) {
        double deltaAngle = abs(ag.angle - current.angle);
        double maxDeltaAngle = timeInterval * maxAngleSpeed;
        if (canTransition(deltaAngle, timeInterval)) {
          // Using the gain as a negative cost
          // TODO 这个0.1设置为参数
          double newCost = current.cost - (0.1 + deltaAngle / maxDeltaAngle) * ag.gain;
          std::pair<double, double> nextNode = make_pair(nextTime, ag.angle);
          if (!dist.count(nextNode) || newCost < dist[nextNode]) {
            dist[nextNode] = newCost;
            prev[nextNode] = {current.time, current.angle};
            pq.push(Node(nextTime, ag.angle, newCost));
          }
        }
      }
    }

    cerr << "No path found within the given constraints." << endl;
    return false;
  }

  // Retrieve the optimal time-angle sequence
  const vector<pair<double, double>> &getOptimalSequence() const
  {
    return optimalSequence;
  }
};

#endif